import os


# 设置 PYARROW_IGNORE_TIMEZONE 环境变量
os.environ['PYARROW_IGNORE_TIMEZONE'] = '1'
import featuretools as ft
# 加载示例数据集
data = ft.demo.load_mock_customer()

# 获取数据框
customers_df = data["customers"]
sessions_df = data["sessions"]
transactions_df = data["transactions"]

# 创建实体集
es = ft.EntitySet(id="customer_data")

# 添加数据框到实体集
es.add_dataframe(dataframe=customers_df,
                 dataframe_name="customers",
                 index="customer_id",
                 time_index="join_date")

es.add_dataframe(dataframe=sessions_df,
                 dataframe_name="sessions",
                 index="session_id",
                 time_index="session_start")

es.add_dataframe(dataframe=transactions_df,
                 dataframe_name="transactions",
                 index="transaction_id",
                 time_index="transaction_time")

# 定义实体之间的关系
relationship_customer_sessions = ft.Relationship(es["customers"]["customer_id"], es["sessions"]["customer_id"])
relationship_sessions_transactions = ft.Relationship(es["sessions"]["session_id"], es["transactions"]["session_id"])

# 将关系添加到实体集
es.add_relationship(relationship_customer_sessions)
es.add_relationship(relationship_sessions_transactions)

# 创建衍生特征
features, feature_names = ft.dfs(entityset=es, target_entity="customers", agg_primitives=["count"],
                                  trans_primitives=["month", "weekday", "day"],
                                  verbose=True)

# 打印生成的特征
print(features.head())
